import time
from collections import defaultdict
from typing import Dict
from flask import request
import threading

class RateLimiter:
    """API速率限制器"""
    
    def __init__(self, max_requests_per_hour: int = 100):
        """
        初始化速率限制器
        
        Args:
            max_requests_per_hour: 每小时最大请求数
        """
        self.max_requests_per_hour = max_requests_per_hour
        self.requests: Dict[str, list] = defaultdict(list)
        self.lock = threading.Lock()
    
    def is_allowed(self, client_id: str = None) -> tuple[bool, int]:
        """
        检查是否允许请求
        
        Args:
            client_id: 客户端ID，如果为None则使用IP地址
            
        Returns:
            (是否允许, 剩余请求数)
        """
        if client_id is None:
            client_id = self._get_client_id()
        
        current_time = time.time()
        hour_ago = current_time - 3600  # 一小时前
        
        with self.lock:
            # 清理过期的请求记录
            self.requests[client_id] = [
                req_time for req_time in self.requests[client_id] 
                if req_time > hour_ago
            ]
            
            # 检查是否超过限制
            if len(self.requests[client_id]) >= self.max_requests_per_hour:
                remaining = 0
                return False, remaining
            
            # 记录当前请求
            self.requests[client_id].append(current_time)
            remaining = self.max_requests_per_hour - len(self.requests[client_id])
            
            return True, remaining
    
    def get_reset_time(self, client_id: str = None) -> int:
        """
        获取速率限制重置时间
        
        Args:
            client_id: 客户端ID
            
        Returns:
            重置时间戳
        """
        if client_id is None:
            client_id = self._get_client_id()
        
        with self.lock:
            if not self.requests[client_id]:
                return int(time.time())
            
            # 返回最早请求的一小时后
            earliest_request = min(self.requests[client_id])
            return int(earliest_request + 3600)
    
    def _get_client_id(self) -> str:
        """获取客户端ID（使用IP地址）"""
        return request.environ.get('HTTP_X_FORWARDED_FOR', request.remote_addr)
    
    def cleanup_expired_records(self):
        """清理过期的请求记录"""
        current_time = time.time()
        hour_ago = current_time - 3600
        
        with self.lock:
            for client_id in list(self.requests.keys()):
                self.requests[client_id] = [
                    req_time for req_time in self.requests[client_id] 
                    if req_time > hour_ago
                ]
                
                # 如果没有记录，删除该客户端
                if not self.requests[client_id]:
                    del self.requests[client_id]
